from pylab import ones, zeros, plot, show, zeros_like, size, shape
import numpy
class airport:
    """This is the class defining airport which has following parameters:
    
    id   <type:int>   : Unique id of the airport
    code <type:str>   : Airport code
    name <type:str>   : Airport name
    ndt  <type:int>   : Number of time periods
    D    <type:int>   : Departure capacity of the airport (Default 50)
    R    <type:int>   : Arrival capacity of the airport (Default 20)
    
    #######################################################

    Functions
    -----------

    set_arrival_capacity(R)    Set the arrival capacity
    set_departure_capacity(D)  Set the departure capacity
    
    
    """
    def __init__(self, id, code, name, ndt, D=50, R=20):
        """Initialise class airport with the values of D(default=50) and R(default=10)"""
        self.id = id
        self.code = code
        self.name = name
        self.D = ones(ndt)*D
        self.R = ones(ndt)*R
    
    def set_arrival_capacity(self, R):
        self.R = R
    
    def set_departure_capacity(self, D):
        self.D = D


class flight:
    """This is the class defining flight which has following parameters:
    
    id           <type:int>      : Unique id of the flight
    code         <type:str>      : Unique code assigned to flight
    k_d          <type:airport>  : Airport from which it is scheduled to depart
    k_a          <type:airport>  : Airport to which it is scheduled to arrive
    d            <type:int>      : Scheduled departure time
    r            <type:int>      : Scheduled arrival time
    c_g          <type:float>    : Ground delay cost
    c_a          <type:float>    : Air delay cost
    next_flight  <type:flight>   : Linked to next flight
    prev_flight  <type:flight>   : Linked to prev flight

    Functions
    -----------
    _set_departure(k_d,d=0)  Set departure airport and time
    _set_arrival(k_a,r)      Set arrival airport and time
    _set_cost(c_g,c_a)       Set the cost of ground and airborne delay
    
    _set_next_flight(flight) Set the next flight
    _set_prev_flight(flight) Set the prev flight

    setup(k_d,d,k_a,r,c_g,c_a) setup departure, arrival and cost together

    Note: If you add a next flight b to flight a make sure to add a as prev flight of b

    """
    def __init__(self, id, code):
        self.id = id
        self.code = code
        self.next_flight = None
        self.prev_flight = None

    def _set_departure(self, k_d, d=0):
        self.k_d = k_d
        self.d = d
    def _set_arrival(self, k_a, r=0):
        self.k_a = k_a
        self.r = r
    def _set_cost(self, c_g=100.0, c_a=100.0):
        self.c_g = c_g
        self.c_a = c_a
    def _set_next_flight(self, flight):
        self.next_flight = flight
    def _set_prev_flight(self, flight):
        self.prev_flight = flight
        
    def setup(self,k_d,d,k_a,r,c_g,c_a):
        self._set_departure(k_d,d)
        self._set_arrival(k_a,r)
        self._set_cost(c_g,c_a)
        
    
class Solver:
    """
    This class is the main Solver which has the following parameters:
    
    Input
    ----------

    ndt        <type:int>   : Number of time periods
    dt         <type:int>   : Time period (default 15 min)
    Airports   <type:list>  : List of all the Airports
    Flights    <type:list>  : List of all the flights
    A          <type:int>   : Upper limit on airborne delay (default 2 time period)
    G          <type:int>   : Upper limit on ground delay (default 4 time period)


    Functions
    ----------
    _setup()               : Setup everything
    
    objective()            : Returns the objective function
    c_airport_capacity()   :Airport Capacity constraint
    c_flight_()            : Flight constraint
    c_on_a()               : Contraint on a
    c_connecting_flight()  : Connecting flight constraint
    
    save_data()            : Dump all the data into an npz file
    
    """
    def __init__(self, ndt, Airports, Flights, A=2, G=4, dt=15, path='../data/data'):
        self.ndt = ndt
        self.dt = dt
        self.Airports = Airports
        self.Flights = Flights
        self.A = A
        self.G = G
        self.path = path
        self.count = 0
    
    def _setup(self):
        """  
        This function sets up all parameters required to solve the problem, which includes:
        u,v,a,g,tau_d,tau_a,airport_d_flight_list,airport_a_flight_list
        -------------
        u <type: 2D array, shape: len(flights) x ndt>
        u[i,j] = 1 if flight i is allowed to depart at time j
        = 0 otherwise
        => sum over j for a particular i is always unity
        
        v <type: 2D array, shape: len(flights) x ndt>
        v[i,j] = 1 if flight i is allowed to land at time j
        = 0 otherwise
        => sum over j for a particular i is always unity
        
        --------------
    
        g <type: 1D array> 
        g[i] contains Number of time period by which a flight is delayed for departure
        
        a <type: 1D array> 
        a[i] contains Number of time period by which a flight is delayed for arrival
    
        ---------------
        
        tau_d<type: 2D array, shape: len(Flights) x (G + 1)>
        tau_d[i,:] or tau_d[i] gives all the possible values of time period when a flight 
        'i' can depart depending upon the upper limit on delays
        
        tau_a<type: 2D array, shape: len(Flights) x (G + A + 1)>
        tau_a[i,:] or tau_a[i] gives all the possible values of time period when a flight 
        'i' can arrive depending upon the upper limit on delays
        
        ---------------
        
        airport_d_flight_list <type: list>
        airport_d_flight_list[i] gives all the flight departing from aiport i
        
        airport_a_flight_list <type: list>
        airport_a_flight_list[i] gives all the flight arriving at aiport i
        
        ---------------
    
        
        """
        
        Flights = self.Flights
        ndt = self.ndt
        # Initialising decision variables

        # u[f,t] = 1 if flight f is assigned to takeoff at time t and 0 otherwise  
        self.u = zeros([len(Flights),2*ndt])

        # v[f,t] = 1 if flight f is assigned to land at time t and 0 otherwise  
        self.v = zeros([len(Flights),2*ndt])
        
        for f in range(len(Flights)):
            self.u[f,Flights[f].d] = 1
            self.v[f,Flights[f].r] = 1
        

        # a[f] and g[f] stands for airborne and ground delay for flight f 
        self.g = ones(len(Flights))
        self.a = ones(len(Flights))
        
        # s is the slack time between two consecutive flights to be performed by the same aircraft
        # s[f] = t[fnext] - t[f]
        self.s = zeros(len(Flights))
        
        for f in self.Flights:
            if f.next_flight != None:
                self.s[f.id] = f.next_flight.d - f.r 
        
        # tau_d and tau_a are the allowable time period on which a flight can land or takeoff
        # this comes from the upper limit on delays
        self.tau_d = zeros([len(Flights), self.G + 1])
        self.tau_a = zeros([len(Flights), self.A + self.G + 1])
       
        ### Need to sort flight by airports which will be used in the capacity constraint
        self.airport_d_flight_list = []
        self.airport_a_flight_list = []
        for i in self.Airports:
            tmpa = []
            tmpd = []
            for f in Flights:
                if f.k_a.id == i.id:
                    tmpa.append(f)
                if f.k_d.id == i.id:
                    tmpd.append(f)
            self.airport_a_flight_list.append(tmpa)
            self.airport_d_flight_list.append(tmpd)
        
        ### Assigning values to tau
        for f in range(len(self.Flights)):
            countd = 0
            countr = 0
            for i in range(2*ndt):
                if i >= self.Flights[f].d and i <= self.Flights[f].d + self.G:
                    self.tau_d[f,countd] = i
                    countd += 1
                if  i >= self.Flights[f].r and i <= self.Flights[f].r + self.G + self.A:
                    self.tau_a[f,countr] = i
                    countr += 1
        
      
    def objective(self):
        """
        Returns the value of objective function.
        !!! All the parameters should be updated before calling objective function
        """
        self.count += 1
        tmp = 0
        for f in range(len(self.Flights)):
            tmp += self.Flights[f].c_a*abs(self.a[f]) + self.Flights[f].c_g*self.g[f]

        return tmp
    
    def c_airport_capacity(self):
        """
        It takes care of the airport's arrival and departure capacity constraint

        return False if the constraint is violated, also it adds square of the 
        value of constraint function to the penalty at every instance of violation.
        
        output [<type:bool>,<type:float>]
        
        """
        self.airport_con_check = []
        mainflag = True
        ndt = self.ndt
        tmp_penalty = 0
        for i in self.Airports:
            flight_list_a = self.airport_a_flight_list[i.id]
            flight_list_d = self.airport_d_flight_list[i.id]

            for t in range(ndt):
                tmp_a = 0
                tmp_d = 0

                for f in flight_list_a:
                    if f.r + self.g[self.Flights.index(f)] + self.a[self.Flights.index(f)]== t:
                        tmp_a += 1

                for f in flight_list_d:
                    if f.d + self.g[self.Flights.index(f)] == t:
                        tmp_d += 1

                tmp_diff_a = i.R[t] - tmp_a
                tmp_diff_d = i.D[t] - tmp_d

                if tmp_diff_a < 0:
                    self.airport_con_check.append(0)
                    mainflag = False
                    tmp_penalty += tmp_diff_a*tmp_diff_a 
                else:
                    self.airport_con_check.append(1)
                
                if tmp_diff_d < 0:
                    mainflag = False
                    tmp_penalty += tmp_diff_d*tmp_diff_d
                    self.airport_con_check.append(0)
                else:
                    self.airport_con_check.append(1)
       
        return [mainflag, tmp_penalty*10000] 
                
    def c_flight_(self):
        """
        This constraint make sure that a flight has a unique time of departure and arrival
        i.e. sum of u and v are unity for a give flight

        It seems like a redundant constraint because it has been taken care in the code
        itself.
        
        return [<type: bool>, 0]
        """
        mainflag = True
        for f in range(len(self.Flights)):
            tmp_d = 0
            tmp_a = 0

            for i in self.tau_d[f]:
                tmp_d += self.u[f,i]
            
            for i in self.tau_a[f]:
                tmp_a += self.v[f,i]
            tmp_a = int(tmp_a)
            tmp_d = int(tmp_d)

            if tmp_d != 1 or tmp_a != 1:
                mainflag = False
                print self.v[f]
                print self.tau_a[f]
        return [mainflag, 0]

    def c_connecting_flight(self):
        """
        Not using currently
        """
        tmp_penalty = 0
        mainflag = True
        for f in self.Flights:
            if f.prev_flight != None:
                tmp_diff = self.g[f.prev_flight.id] + self.a[f.prev_flight.id]\
                    - (self.s[f.prev_flight.id]) - self.g[f.id] 
                if tmp_diff > 0:
                    mainflag =  False
                    tmp_penalty += tmp_diff*tmp_diff 
                    
       
        return [mainflag, tmp_penalty*10000]
    
    def c_on_a(self):
        """
        
        To make sure that self.a > 0
        
        return [<type: bool>, <type: float>]
        
        """
        tmp_penalty = 0
        mainflag = True
        for f in range(len(self.Flights)):
            if self.a[f] < 0:
                mainflag = False
                tmp_penalty += self.a[f]*self.a[f]

        return [mainflag, tmp_penalty*1000000]
    
    
            
    
    def _update_ag(self,ind_a,ind_g):
        """
        Update a and g to the new values

        input <index of tau_a, index of tau_d>
        """
        for f in range(len(self.Flights)):
            self.g[f] = self.tau_d[f,ind_g[f]] - self.Flights[f].d
            self.a[f] = self.tau_a[f,ind_a[f]] - self.Flights[f].r - self.g[f]
    def _update_uv(self):
        """
        Updates u and v from the updated values of a and g
        Input: None
        
        should be called after updating a and g

        !!!Warning
        Here is a catch if suppose expected time of arrival becomes greater than the
        total defines time period, it will create a problem here because index of v 
        will become greated then its size.
        Solution 
        - define a large size of time period and make sure that p number of 
        last slot are not assigned to any flight during initialization. So that this 
        limit is never reached.
        
        - or define a array which will take the values if index > len - 1, its same 
        like shifting a flight to early hours of day 2
        """
        self.u = zeros_like(self.u)
        self.v = zeros_like(self.v)
        for i in range(len(self.Flights)):
            self.u[i, self.g[i] + self.Flights[i].d] = 1
            self.v[i, self.a[i] + self.g[i] + self.Flights[i].r] = 1
           
    def c_on_ind(self,ind_a,ind_g):
        """
        !!! Not required now just keeping it.

        Constraint on index of a and g 
        This it not a physical Constraint it comes from the method we are using for 
        solving the problem

        """
        mainflag = True
        for f in range(len(self.Flights)):
            if ind_a[f] >= len(self.tau_a[f]) or ind_g[f] >= len(self.tau_d[f]):
                mainflag = False
            
        return mainflag
    
       
    
